#include "stdafx.h"

//#define EXPORT extern "C" __declspec(dllexport)
//typedef TCHAR * PTCHAR;
#ifdef UNICODE
#define T L
#else
#define T
#endif

TCHAR TPath[MAX_PATH]=".";
TCHAR MPath[MAX_PATH]=".";
bool ChCur()
{
	if (*TPath=='.') return true;
	if (!SetCurrentDirectory(TPath)) return false;
	MyPrintf("Chdir: '%s'\n",TPath);
	return true;
}

bool ChReg()
{
	if (*MPath=='.') if (!GetReg()) return false;

	if (!SetCurrentDirectory(MPath)) return false;
	MyPrintf("Chdir: '%s'\n",MPath);
	return true;
}

PTCHAR McKey= "SOFTWARE\\Network Associates\\TVD\\Shared Components\\VirusScan Engine\\4.0.xx";
PTCHAR Frame="SOFTWARE\\Network Associates\\TVD\\Shared Components\\Framework";
PTCHAR Path="Installed Path";// - where Nai placed

DWORD MakePath(char * buf, DWORD size, const char * path, const char * param)
{
	HKEY key;
	bool res=true;
	*buf=0;
	DWORD wr=size;
	if (ERROR_SUCCESS!=RegOpenKeyEx(HKEY_LOCAL_MACHINE,path,0,KEY_QUERY_VALUE,&key))
		{InformError("Open reg",path);return false;}

	if (ERROR_SUCCESS!=RegQueryValueEx(key,param,0,0,(LPBYTE)buf,&wr))
	{InformError("Query key",param);res=false;}
	if (ERROR_SUCCESS!=RegCloseKey(key)) res=false;
	if (res) {wr=lstrlen(buf);while (wr>1 && buf[wr-1]=='\\') buf[--wr]=0;}
	if (!res) return false;
	return wr;
}

bool GetReg()
{
	HKEY key;	
	DWORD wr=sizeof(MPath);
	bool res=true;

	if (*MPath=='.')
	{
	if (ERROR_SUCCESS!=RegOpenKeyEx(HKEY_LOCAL_MACHINE,McKey,0,KEY_QUERY_VALUE,&key))
		{InformError("Open reg",McKey);return false;}

	if (ERROR_SUCCESS!=RegQueryValueEx(key,"DAT",0,0,(LPBYTE)MPath,&wr)) {InformError("Query key",McKey);res=false;}
	if (ERROR_SUCCESS!=RegCloseKey(key)) res=false;
	if (res) {wr=lstrlen(MPath);while (wr>1 && MPath[wr-1]=='\\') MPath[--wr]=0;}
	if (!res) return false;
	}

	if (*TPath == '.')
	{
		wr=GetCurrentDirectory(sizeof(TPath),TPath);
		if (!wr || wr>sizeof(TPath)) {return false;}
	}
	return true;
}

const unsigned rtptr=0x1BB;
bool Upd2Rtp(PTCHAR file, PTCHAR dst=0)
{
	bool res=true;
	FILETIME c,a,w;

	HANDLE in=CreateFile(file,GENERIC_READ,FILE_SHARE_READ,0,OPEN_EXISTING,FILE_ATTRIBUTE_NORMAL|FILE_FLAG_SEQUENTIAL_SCAN,0);
	if (in==INVALID_HANDLE_VALUE) {InformError("Open for read",file);return false;}

	GetFileTime(in,&c,&a,&w);
	unsigned off,len,size=GetFileSize(in,0);

	HANDLE FM=CreateFileMapping(in,0,PAGE_READONLY,0,size,0);
	if (FM!=INVALID_HANDLE_VALUE)
	{
	unsigned char * src=(unsigned char*)MapViewOfFile(FM,FILE_MAP_READ,0,0,0);
	if (src)
	{
	DWORD wr;
	bool IsUPD=false;

	if ((size>rtptr+8) && *(DWORD*)(src+2)=='1DPU')
	{
		off=*(unsigned*)(src+rtptr+0);
		len=*(unsigned*)(src+rtptr+4);

		if (off>=size || len>=size || (off+len)>=size || off<rtptr+8)
			IsUPD=false;
		else IsUPD=true;
	}

	if (IsUPD)
	{

	if (!dst)
	{

	int len=lstrlen(file);
	dst=file+len;
	if (len>4 && !lstrcmpi(dst-4,".upd"))
		dst-=4;
	  lstrcpy(dst,".rtp");

	dst=file;
	
	}
	HANDLE out=CreateFile(dst,GENERIC_READ,FILE_SHARE_READ,0,OPEN_ALWAYS,FILE_ATTRIBUTE_NORMAL,0);
	if (out!=INVALID_HANDLE_VALUE)
	{
		size=GetFileSize(out,0);
		CloseHandle(out);
	}

	if (size != len)
	{

	out=CreateFile(dst,GENERIC_WRITE,FILE_SHARE_READ,0,CREATE_ALWAYS,FILE_ATTRIBUTE_NORMAL|FILE_FLAG_SEQUENTIAL_SCAN,0);
	if (out!=INVALID_HANDLE_VALUE)
	{

	if (!WriteFile(out,src+off,len,&wr,0))
	{InformError("Can not write",dst);res=false;}
	
	SetFileTime(out,&c,&a,&w);
	CloseHandle(out);
	} else {InformError("Open for write",dst);res=false;}
	}
	}

	UnmapViewOfFile(src);
	} else {InformError("Can not map",file);res=false;}
	CloseHandle(FM);
	} else {InformError("Can not create Mapping",file);res=false;}
	CloseHandle(in);
	return res;
}
char EicarString[]="X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*";
bool Eicar(LPCSTR file)
{
	bool res=true;
	HANDLE out=CreateFile(file,GENERIC_WRITE,FILE_SHARE_READ,0,CREATE_ALWAYS,FILE_ATTRIBUTE_NORMAL|FILE_FLAG_SEQUENTIAL_SCAN,0);
	if (out!=INVALID_HANDLE_VALUE)
	{
		DWORD wr;
	if (!WriteFile(out,EicarString,sizeof(EicarString)-1,&wr,0))
	{InformError("Can not write",file);res=false;}
	CloseHandle(out);
	} else {InformError("Open for write",file);res=false;}
	return res;
}

void MyPrintf(LPCSTR fmt, ...)
{
	va_list va;
	va_start(va,fmt);
	char tmp[4*1024];
	DWORD wr;
	WriteFile(GetStdHandle(STD_OUTPUT_HANDLE),tmp,wvsprintf(tmp,fmt,va),&wr,0);
	va_end(va);
}

int InformError(LPCSTR msg, LPCSTR arg, UINT show)
{
	MyPrintf("%s:%s\n",msg,arg);
if (show)
{
LPVOID lpMsgBuf;
FormatMessage( 
    FORMAT_MESSAGE_ALLOCATE_BUFFER | 
    FORMAT_MESSAGE_FROM_SYSTEM | 
    FORMAT_MESSAGE_IGNORE_INSERTS,
    NULL,GetLastError(),
    0,(LPTSTR) &lpMsgBuf,0,NULL);
	
	MyPrintf("%s\n",lpMsgBuf);

//	int res=MessageBox(NULL,(LPCTSTR)lpMsgBuf,Args,show);
	LocalFree( lpMsgBuf );
}
	return IDOK;
}
UINT (__cdecl * Decrypt2)(UINT,LPVOID,UINT,LPVOID,DWORD*)=0;
const CHAR nai[]="NaiSign.Dll";
bool Unpack2(PTCHAR file)
{
	bool res=true;
	PTCHAR dst=0;
	if (!Decrypt2)
	{
	char naipath[270];
	if (MakePath(naipath,sizeof(naipath),Frame,Path))
		lstrcat(naipath,"\\");
	lstrcat(naipath,nai);

	HMODULE mod=LoadLibrary(naipath);
	(*(FARPROC*)&Decrypt2) = GetProcAddress(mod,"naisign_Reserved");
	if (!Decrypt2) {InformError("Can not load library",naipath);return false;}
	}

	FILETIME c,a,w;

	HANDLE in=CreateFile(file,GENERIC_READ,FILE_SHARE_READ,0,OPEN_EXISTING,FILE_ATTRIBUTE_NORMAL|FILE_FLAG_SEQUENTIAL_SCAN,0);
	if (in==INVALID_HANDLE_VALUE) {InformError("Open for read",file);return false;}

	GetFileTime(in,&c,&a,&w);
	unsigned off,len,size=GetFileSize(in,0);

	HANDLE FM=CreateFileMapping(in,0,PAGE_READONLY,0,size,0);
	if (FM!=INVALID_HANDLE_VALUE)
	{
	unsigned char * src=(unsigned char*)MapViewOfFile(FM,FILE_MAP_READ,0,0,0);
	if (src)
	{
	DWORD wr;

	if (!dst)
	{
		dst=file;
		while (*dst && *dst!='.') dst++;
		if (!*dst) {*file='~';dst=file;}
		else
		{
			if (!lstrcmpi(dst,".McS")) lstrcpy(dst,".mc");
			else if (!lstrcmpi(dst,".z")) lstrcpy(dst,".u");
			else lstrcpy(dst,"~");
			dst=file;
		}
	}

	HANDLE out=CreateFile(dst,GENERIC_WRITE,FILE_SHARE_READ,0,CREATE_ALWAYS,FILE_ATTRIBUTE_NORMAL|FILE_FLAG_SEQUENTIAL_SCAN,0);
	if (out!=INVALID_HANDLE_VALUE)
	{
		LPVOID buf=VirtualAlloc(NULL,size,MEM_COMMIT,PAGE_READWRITE);
		// check: Starting with "EPO"
		wr=size;

		Decrypt2(1,src,size,buf,&wr);
		
	if (!WriteFile(out,buf,wr,&wr,0))
	{InformError("Can not write",dst);res=false;}

	VirtualFree(buf,0,MEM_RELEASE);
//	FreeLibrary(mod);
	
	SetFileTime(out,&c,&a,&w);
	CloseHandle(out);
	} else {InformError("Open for write",dst);res=false;}
	}

	UnmapViewOfFile(src);
	} else {InformError("Can not map",file);res=false;}
	CloseHandle(FM);
	CloseHandle(in);
	return res;
}

bool __cdecl DecodeFile(const PTCHAR file, bool encode)
{
	bool res=true;
	FILETIME c,a,w;

	HANDLE out=CreateFile(file,GENERIC_READ|GENERIC_WRITE,FILE_SHARE_READ|FILE_SHARE_WRITE,0,OPEN_EXISTING,FILE_ATTRIBUTE_NORMAL|FILE_FLAG_SEQUENTIAL_SCAN,0);
	if (out==INVALID_HANDLE_VALUE) {InformError("Open for read-write",file);return false;}

	GetFileTime(out,&c,&a,&w);

	unsigned size=GetFileSize(out,0);
//	SetFilePointer(out,0,0,0);

	HANDLE FMout=CreateFileMapping(out,0,PAGE_READWRITE,0,size,0);	
	if (FMout != INVALID_HANDLE_VALUE)
	{
	unsigned char * bout=(unsigned char*)MapViewOfFile(FMout,FILE_MAP_READ|FILE_MAP_WRITE,0,0,0);
	if (bout)
	{
	res=DoMcDecode(bout,size,encode);

	UnmapViewOfFile(bout);
	} else InformError("Do map",file);
	CloseHandle(FMout);
	} else InformError("Create Mapping",file);

//	SetFilePointer(out,size,0,0);
//	SetEndOfFile(out);

	SetFileTime(out,&c,&a,&w);
	CloseHandle(out);
	MyPrintf("%s - %s %s\n",file,encode?"Encode":"Decode",res?"OK":"Fail");
	return res;
}

bool __cdecl GetVer(const PTCHAR file, bool nothing)
{
	HANDLE in=CreateFile(file,GENERIC_READ,FILE_SHARE_READ,0,OPEN_ALWAYS,FILE_ATTRIBUTE_NORMAL,0);
	if (in==INVALID_HANDLE_VALUE) {InformError("Open for read",file);return false;}
	
	VerInfo vi;
	DWORD wr;

	ReadFile(in,&vi,sizeof(vi),&wr,NULL);
	CloseHandle(in);

	if (memcmp(vi.sign,DAT_CopyRight,sizeof(DAT_CopyRight))) return false;

	bool Packed=true;
	if (vi.pck != DAT_PACKED)
	{
	unsigned char enc[256];
	MakeTable(enc, vi.salt);
	DoDecode (enc, (unsigned char*)&vi, VerInfo_Plain, sizeof(vi));
	Packed=false;
	}
	
	MyPrintf("%04u-%02u-%02u\t%s\t%u\t%u\t%s\n",(vi.day >> 9) + 1980, (vi.day >> 5) & 15, vi.day & 31, Packed?"PCK":vi.pck != DAT_PACKED?"BAD":"UNP", vi.ver, vi.vir1 + vi.vir2, file);

	return true;
}

PTCHAR RARargs="e -av- -u";
bool __cdecl UnrarFile(const PTCHAR file, bool encode)
{
	PROCESS_INFORMATION pi;
	STARTUPINFO si;

	char cmd[MAX_PATH*3];
	memset(cmd,0,sizeof(cmd));
	wsprintf(cmd,"\"rar.exe\" %s \"%s\" \"%s\\\"",RARargs,file,MPath);
//	wsprintf(cmd,"e -av- -u \"%s\" \"",file);lstrcat(cmd,MPath);lstrcat(cmd,"\\\"");

	GetStartupInfo(&si);
	MyPrintf("Executing: '%s'\n",cmd);
	if (!CreateProcess(NULL,cmd,NULL,NULL,false,0,NULL,NULL,&si,&pi))
	{InformError("Execute failed",cmd);return false;}
	CloseHandle(pi.hThread);
	WaitForSingleObject(pi.hProcess,10000);
	CloseHandle(pi.hProcess);
	return true;
}
PTCHAR Library="PatchW32.dll";

UINT (CALLBACK *RTPatchApply32)( LPSTR CmdLine, LPVOID (CALLBACK * CallBackFn)(UINT, LPVOID), BOOL WaitFlag)=0;
bool __cdecl ApplyPatch(const PTCHAR file, bool encode)
{
	char newFile[1024];
	char cmd[1024];

	lstrcpyn(newFile,file,sizeof(newFile));
	Upd2Rtp(newFile);

	if (!RTPatchApply32)
	{
	HMODULE hMod=LoadLibrary(Library);
	if (hMod)
	*(FARPROC*)&RTPatchApply32=GetProcAddress(hMod,"RTPatchApply32@12");
	}

	if (!RTPatchApply32) {InformError("Load library",Library);return false;}
	wsprintf(cmd,"\"%s\" %s",MPath,newFile);
	MyPrintf("Patching: '%s'\n",cmd);
	return 0!=RTPatchApply32(cmd,PatchCallBack,1);
}
PTCHAR service="McShield";
PTCHAR machine=NULL;
SC_HANDLE sm,sh;

void SVClose()
{
	if (sh)CloseServiceHandle(sh),sh=NULL;
	if (sm)CloseServiceHandle(sm),sm=NULL;
}

bool ServiceControl(bool stop)
{
	bool res=true;
	SERVICE_STATUS ss;
	if (!sm) sm=OpenSCManager((machine && *machine)?machine:NULL,NULL, SC_MANAGER_CONNECT|
		SERVICE_START|SERVICE_STOP|SERVICE_QUERY_STATUS);
	if (!sm) {InformError("Open serivce manager",machine);return false;}
	if (!sh)
	sh=OpenService(sm,service,SERVICE_START|SERVICE_STOP|SERVICE_QUERY_STATUS);
	if (!sh) {InformError("No service",service);return false;}
	else
	{
		if (!stop) {if (!StartService(sh,0,NULL)) res=false;}
		else {if (!ControlService(sh,SERVICE_CONTROL_STOP,&ss)) res=false;}

		if (res)
		{
			if (!stop)
			if(!QueryServiceStatus(sh,&ss)){InformError("Service query",service);res=false;}
//			DWORD dwStartTickCount = GetTickCount();

			DWORD dwWaitTime = ss.dwWaitHint / 10;
			if( dwWaitTime < 1000 ) dwWaitTime = 1000;
			else if ( dwWaitTime > 3000 ) dwWaitTime = 3000;

			for (int run=3;run>0 && res;run--) 
			{
				switch (ss.dwCurrentState)
				{
				case SERVICE_START_PENDING:break;
				case SERVICE_STOP_PENDING:break;
				default:run=0;
				}
				if (run<=0) break;

				Sleep( dwWaitTime );
				if (!QueryServiceStatus(sh,&ss)) {InformError("Service query",service);res=false;}
			}
		} else InformError("Service control",service);
	}
	return res;
}

bool WINAPI cAllFiles(const PTCHAR name, bool(__cdecl* Func)(const PTCHAR,bool) , bool arg)
{
		char CurFile[MAX_PATH];
		bool local=*name=='.' && name[1]=='\\';
		if (local) wsprintf(CurFile,"%s%s",TPath,name+1);
		else lstrcpy(CurFile,name);

		if (strchr(name,'?')||strchr(name,'*'))
		{
		bool res=true;
		WIN32_FIND_DATA fd;
		HANDLE ff=FindFirstFile(CurFile,&fd);
		if (ff==INVALID_HANDLE_VALUE) {InformError("Not Found",name);return false;}
		do{
			if (local) wsprintf(CurFile,"%s\\%s",TPath,fd.cFileName);
			else lstrcpy(CurFile,fd.cFileName);
			if (!Func(CurFile,arg)) res=false;
		}while(FindNextFile(ff,&fd));
		FindClose(ff);
		return res;
		} else return Func(CurFile,arg);
}

void Help()
{
	MyPrintf("McAfee control tool v1.4 by Gimly\n\nAllowed keys:\n"
	"\t/ver\t- Get DAT files version. Note: PCK is Only format accepted by AV.\n"
	"\t/reg\t- Get registry Engine path for DAT operations.\n"
	"\t/chreg\t- change current dir to Engine dir.\n"
	"\t/cur\t- change current dir back to startup.\n"
	"\t/decode\t- Decode DAT files. (Prepare for patching).\n"
	"\t/encode\t- Encode DAT files. (AVirus knows only that format).\n"
	"\t/patch\t- Apply UPD or RTP path to DAT.\n"
	"\t/unrar\t- Unrar archive of DAT files.\n"
	"\t/unp2\t- Decrypt .McS (EPO) files.\n"
	"\n"
	"\t/stop\t- Stop McShield service.\n"
	"\t/start\t- Start McShield service.\n"
	"\t/svnew\t- Close Service Control Manager (to change /serv or /mach).\n"
	"\t/serv\t- Change service name instead '%s'. Possible AvSynMgr for v7pro.\n"
	"\t/mach\t- Change machine name instead local for service control.\n"
	"\t/args\t- Unrar arguments instead '%s'.\n"
	"\t/help\t- This help\n"
	"\t/eicar\t- Make EiCar test virus file\n"	
	"\n"
,service,RARargs);
}
const PTCHAR AllDat=T"*.dat";

int main(int argc, TCHAR* argv[])
{
	
//	int argc=0;PTCHAR * argv=CommandLineToArgv(GetCommandLine(), &argc);
	for (int i=1;i<argc;i++)
	{
		bool enc;	

		if (!lstrcmpi(argv[i],"/ver"))
		{
				MyPrintf("Date\t\tPck\tVer\tCount\tName\n");
				cAllFiles(AllDat,GetVer,false);
		}
		else if ((enc=false,!lstrcmpi(argv[i],"/decode")) ||
				  (enc=true,!lstrcmpi(argv[i],"/encode")))
			cAllFiles(AllDat,DecodeFile,enc);
		else if (!lstrcmpi(argv[i],"/reg"))		GetReg();
		else if (!lstrcmpi(argv[i],"/chReg"))	ChReg();
		else if (!lstrcmpi(argv[i],"/cur"))		ChCur();
		else if ((enc=true,!lstrcmpi(argv[i],"/stop"))||
				(enc=false,!lstrcmpi(argv[i],"/start")))
			ServiceControl(enc);
		else if (!lstrcmpi(argv[i],"/svnew"))	SVClose();
		else if (!lstrcmpi(argv[i],"/help")||!lstrcmpi(argv[i],"/?")) Help();
		else
		{
			if (argc<=i+1) {InformError("Ignored",argv[i],0);continue;}

			 if (!lstrcmpi(argv[i],"/patch")) cAllFiles(argv[++i],ApplyPatch,false);
		else if (!lstrcmpi(argv[i],"/unrar")) cAllFiles(argv[++i],UnrarFile,false);
		else if (!lstrcmpi(argv[i],"/unp2")) Unpack2(argv[++i]);
		else if (!lstrcmpi(argv[i],"/args")) RARargs=argv[++i];
		else if (!lstrcmpi(argv[i],"/serv")) service=argv[++i];
		else if (!lstrcmpi(argv[i],"/mach")) machine=argv[++i];
		else if (!lstrcmpi(argv[i],"/eicar")) Eicar(argv[++i]);
		else InformError("Ignored",argv[i],0);
		}
	}
	ExitProcess(0);
	return 0;
}
